## #
## 存储模块
## 存储所有的代理数据
## #

from mimetypes import init
import redis
from proxypool.exceptions import PoolEmptyException
from proxypool.schemas.proxy import Proxy
from proxypool.settings import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD,REDIS_DB,REDIS_CONNECTION_STRING,\
    PROXY_SCORE_MAX, PROXY_SCORE_MIN, REDIS_KEY, PROXY_NUMBER_MAX, PROXY_SCORE_MIN, PROXY_SCORE_INIT

from random import choice
from typing import List
from loguru import logger
from proxypool.utils.proxy import is_valid_proxy, convert_proxy_or_proxies

REDIS_CLIENT_VERSION = redis.__version__
IS_REDIS_VERSION_2 = REDIS_CLIENT_VERSION.startswith('2.')

class StorageClient(object):
    '''
    存储客户端
    '''
    def __init__(self, host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB,
                 connection_string=REDIS_CONNECTION_STRING, **kwargs):
        """
        如果连接不存在则使用host
        """
        if connection_string:
            self.db = redis.StrictRedis.from_url(connection_string, decode_responses=True, **kwargs)
        else:
            self.db = redis.StrictRedis(
                host=host, port=port, password=password, db=db, decode_responses=True, **kwargs)
    
    def add(self, proxy:Proxy, score=PROXY_SCORE_INIT) -> init:
        """
        添加代理，并且设置分数
        """
        if not is_valid_proxy(proxy.string()):
            logger.info(f'失效代理 {proxy}, 直接返回')
            return
        if not self.exists(proxy):
            if IS_REDIS_VERSION_2:
                return self.db.zadd(REDIS_KEY, score, proxy.string())
            return self.db.zadd(REDIS_KEY, {proxy.string(): score})
      
    def exists(self, proxy:Proxy):
        # 如果hash中file是否存在
        return self.db.zscore(REDIS_KEY, proxy.string()) is not None

    def random(self) -> Proxy:
        """
        获取随机代理获取
        """
        proxies = self.db.zrangebyscore(REDIS_KEY, PROXY_SCORE_MAX, PROXY_SCORE_MAX)
        if len(proxies):
            return convert_proxy_or_proxies(choice(proxies))
        
        proxies = self.db.zrangebyscore(REDIS_KEY, PROXY_SCORE_MIN, PROXY_SCORE_MAX)
        if len(proxies):
            return convert_proxy_or_proxies(choice(proxies))
        
        raise PoolEmptyException

    def decrease(self, proxy: Proxy) -> int:
        '''
        代理扣减
        '''
        # 分数扣减1
        if IS_REDIS_VERSION_2:
            self.db.zincrby(REDIS_KEY, proxy.string(), -1)
        else:
            self.db.zincrby(REDIS_KEY, -1, proxy.string())
        # 获取分数
        score = self.db.zscore(REDIS_KEY, proxy.string())
        logger.info(f'{proxy.string()} score -1, 当前 {score}')
        # 校验是否小于等于最小值
        if score <= PROXY_SCORE_MIN:
            logger.info(f'{proxy.string()} 当前 score {score}，小于最小值, 移除')
            self.db.zrem(REDIS_KEY, proxy.string())

    def max(self, proxy: Proxy) -> int:
        """
        设置代理分数最大值
        """
        logger.info(f'{proxy.string()} 有效, 设置 {PROXY_SCORE_MAX}')
        if IS_REDIS_VERSION_2:
            return self.db.zadd(REDIS_KEY, PROXY_SCORE_MAX, proxy.string())
        return self.db.zadd(REDIS_KEY, {proxy.string(): PROXY_SCORE_MAX})

    def count(self) -> int:
        """
        计算当前代理数
        """
        return self.db.zcard(REDIS_KEY)

    def all(self) -> List[Proxy]:
        """
        获取所有代理
        """
        return convert_proxy_or_proxies(self.db.zrangebyscore(REDIS_KEY, PROXY_SCORE_MIN, PROXY_SCORE_MAX))

    def batch(self, cursor, count) -> List[Proxy]:
        """
        获取批量代理
        """
        cursor, proxies = self.db.zscan(REDIS_KEY, cursor, count=count)
        return cursor, convert_proxy_or_proxies([i[0] for i in proxies])

    def get_db(self) -> redis.Redis:
        return self.db

